#!/bin/bash

set -euo pipefail

# This script ensures that the EVM contracts can be safely upgraded to without
# bricking the contracts. It does this by simulating contract upgrades against
# the mainnet state, and checks that the state is consistent after the upgrade.
#
# By default, the script will compile the contracts and run the upgrade. It's
# possible to simulate an upgrade against an already deployed implementation
# contract (which is useful for independent verification of a governance
# proposal) -- see the usage instructions below.

function usage() {
cat <<EOF >&2
Usage:

  $(basename "$0") [-h] [-m s] [-c s] [-x] [-k] [-d] [-a s] [-l s] [-s] -- Simulate an upgrade on a fork of mainnet, and check for any errors.

  where:
    -h  show this help text
    -m  module (bridge, token_bridge, nft_bridge)
    -c  chain name
    -x  run anvil
    -d  don't compile contract first
    -k  keep anvil alive
    -l  file to log to (by default creates a new tmp file)
    -a  new code address (by default it builds the most recent contract in the repository)
    -s  shutdown
EOF
exit 1
}

before=$(mktemp)
after=$(mktemp)

### Parse command line options
address=""
module=""
chain_name=""
run_anvil=false
skip_compile=false
keepalive_anvil=false
shutdown=false
anvil_out=$(mktemp)
while getopts ':hm:c:a:xkdl:s' option; do
  case "$option" in
    h) usage
       ;;
    m) module=$OPTARG
       ;;
    a) address=$OPTARG
       ;;
    c) chain_name=$OPTARG
       ;;
    x) run_anvil=true
       ;;
    d) skip_compile=true
       ;;
    l) anvil_out=$OPTARG
       ;;
    k) keepalive_anvil=true
       run_anvil=true
       ;;
    s) shutdown=true
       ;;
    :) printf "missing argument for -%s\n" "$OPTARG" >&2
       usage
       ;;
   \?) printf "illegal option: -%s\n" "$OPTARG" >&2
       usage
       ;;
  esac
done
shift $((OPTIND - 1))

# Check that we have the required arguments
[ -z "$chain_name" ] && usage
[ -z "$module" ] && usage

# Get core contract address
CORE=$(worm info contract mainnet "$chain_name" Core)
echo "core: $CORE"

# Use the local devnet guardian key (this is not a production key)
GUARDIAN_ADDRESS=0xbeFA429d57cD18b7F8A4d91A2da9AB4AF05d0FBe
GUARDIAN_SECRET=cfb12303a19cde580bb4dd771639b0d26bc68353645571a8cff516ab2ee113a0

ANVIL_PID=""

function clean_up () {
    ARG=$?
    [ -n "$ANVIL_PID" ] && kill "$ANVIL_PID"
    exit $ARG
}
trap clean_up SIGINT SIGTERM EXIT


#TODO: make RPC an optional argument
HOST="http://0.0.0.0"
PORT="8545"
RPC="$HOST:$PORT"

if [[ $run_anvil = true ]]; then
    ./anvil_fork "$chain_name"
    ANVIL_PID=$!
    echo "🍴 Forking mainnet..."
    echo "Anvil logs in $anvil_out"
    sleep 5
    # ps | grep "$ANVIL_PID"
fi

MODULE=""
SCRIPT=""
case "$module" in
    bridge|core)
        MODULE=Core
        if [[ $shutdown = true ]]; then
            SCRIPT="DeployCoreShutdown.s.sol:DeployCoreShutdown"
            SOLFILE="DeployCoreShutdown.s.sol"
        else
            SCRIPT="DeployCoreImplementationOnly.s.sol:DeployCoreImplementationOnly"
            SOLFILE="DeployCoreImplementationOnly.s.sol"
        fi
        ;;
    token_bridge)
        MODULE=TokenBridge
        if [[ $shutdown = true ]]; then
            SCRIPT="DeployTokenBridgeShutdown.s.sol:DeployTokenBridgeShutdown"
            SOLFILE="DeployTokenBridgeShutdown.s.sol"
        else
            SCRIPT="DeployTokenBridgeImplementationOnly.s.sol:DeployTokenBridgeImplementationOnly"
            SOLFILE="DeployTokenBridgeImplementationOnly.s.sol"
        fi
        ;;
    nft_bridge)
        MODULE=NFTBridge
        if [[ $shutdown = true ]]; then
            SCRIPT="DeployNFTBridgeShutdown.s.sol:DeployNFTBridgeShutdown"
            SOLFILE="DeployNFTBridgeShutdown.s.sol"
        else
            SCRIPT="DeployNFTBridgeImplementationOnly.s.sol:DeployNFTBridgeImplementationOnly"
            SOLFILE="DeployNFTBridgeImplementationOnly.s.sol"
        fi
        ;;
    *) echo "unknown module $module" >&2
       usage
       ;;
esac

CONTRACT=$(worm info contract mainnet "$chain_name" "$MODULE")

EVM_CHAIN_ID=$(printf "%d" $(curl http://localhost:8545/ -X POST -H "Content-Type: application/json" --data '{"method":"eth_chainId","params":[],"id":1,"jsonrpc":"2.0"}' -s | jq -r .result))

# Step 1) Figure out the contract address depending on the flags -- either use
# an address passed in as an argument, or use the most recent contract in the repo.
if [[ -n "$address" ]]; then
    new_implementation="$address"
else
    if [[ $skip_compile = false ]]; then
        echo "🛠  Compiling contract..."
        build_output=$(npm run build) || ( echo "$build_output" && exit 1 )
    fi
    printf "⬆️  Deploying implementation..."
    forge script ./forge-scripts/${SCRIPT} \
        --rpc-url "$RPC" \
        --private-key "0x4f3edf983ac636a65a842ce7c78d9aa706d3b113bce9c46f30d7d21715b23b1d" \
        --broadcast \
        --silent
    returnInfo=$(cat ./broadcast/${SOLFILE}/$EVM_CHAIN_ID/run-latest.json)
    # Extract the address values from 'returnInfo'
    new_implementation=$(jq -r '.returns.deployedAddress.value' <<< "$returnInfo")
fi
printf " %s\n" "$new_implementation"

# Step 2) generate upgrade VAA using the local guardian key
vaa=$(worm generate upgrade -c "$chain_name" -a "$new_implementation" -m $MODULE -g "$GUARDIAN_SECRET")

# Step 3) the VAA we just signed in Step 2) is not compatible with the guardian
# set on mainnet (since that corresponds to a mainnet guardian network). We need
# to thus locally replace the guardian set with the local guardian key.
echo "💂 Overriding guardian set with $GUARDIAN_ADDRESS"
worm evm hijack -g "$GUARDIAN_ADDRESS" -i 0 -a "$CORE" --rpc "$RPC"> /dev/null

# Step 4) query state before upgrade
echo "🔍 Querying old contract state"
worm evm info -c "$chain_name" -m $MODULE -n devnet -a "$CONTRACT" --rpc "$RPC" | grep -v '"implementation":' > "$before"

# Step 5) upgrade contract
echo "🤝 Submitting VAA"
worm submit "$vaa" -n devnet -a "$CONTRACT" --rpc "$RPC" > /dev/null

# Step 6) query state after upgrade
echo "🔍 Querying new contract state"
worm evm info -c "$chain_name" -m $MODULE -n devnet -a "$CONTRACT" --rpc "$RPC" | grep -v '"implementation":' > "$after"

# Step 7) compare old and new state and exit with error if they differ
git diff --no-index "$before" "$after" --exit-code && echo "✅ Upgrade simulation successful" || exit 1

# Anvil can be kept alive by setting the -k flag. This is useful for interacting
# with the contract after it has been upgraded.
if [[ $keepalive_anvil = true ]]; then
    echo "Listening on $RPC"
    # tail -f "$anvil_out"
    wait "$ANVIL_PID"
fi
